from paddleocr import PaddleOCR, draw_ocr
from PIL import Image
from common import IdCardStraight
import os
import glob


class OCR:
    def __init__(self):
        self.ocr = PaddleOCR(use_angle_cls=True, lang="ch")

    # need to run only once to download and load model into memory
    def starOCR(self, img_path):
        result = self.ocr.ocr(img_path, cls=True)
        # image = Image.open(img_path).convert('RGB')
        boxes = [line[0] for line in result[0]]
        txts = [line[1][0] for line in result[0]]
        scores = [line[1][1] for line in result[0]]
        # im_show = draw_ocr(image, boxes, txts, scores)
        # im_show = Image.fromarray(im_show)
        print(f"OCR结果 {txts}")
        return txts
        # im_show.save('result.jpg')
        # im_show.show()

    def parseIdCard(self, txts):
        postprocessing = IdCardStraight(txts)
        id_result = postprocessing.run()
        return id_result

    # png_files = glob.glob(os.path.join("F:\\train_data", '*.jpg'))
    #
    # for file in png_files:
    #     starOCR(file)

    def run(self):
        txts = self.starOCR('testImage/testIdCard.png')
        print(self.parseIdCard(txts))
        input('')


if __name__ == '__main__':
    try:
        test = OCR()
        test.run()
    except Exception as e:
        print(e)
